"""
示例 5: 修改现有 PA 文件

演示如何：
1. 读取现有的 PA 文件
2. 修改电极和电压
3. 添加新的电极结构
4. 保存修改后的文件
"""

import sys
import os

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, project_root)

from SIMION.PA import PA


def modify_electrode_voltage(input_file, output_file):
    """修改电极电压"""
    
    print("\n" + "="*60)
    print("修改电极电压")
    print("="*60 + "\n")
    
    try:
        # 读取现有文件
        print(f"读取文件: {input_file}")
        pa = PA(file=input_file)
        
        print(f"阵列尺寸: {pa.nx()} x {pa.ny()} x {pa.nz()}")
        print(f"原始最大电压: {pa.max_voltage()} V\n")
        
        # 统计原始电极
        electrode_count = 0
        voltage_sum = 0.0
        
        for z in range(pa.nz()):
            for y in range(pa.ny()):
                for x in range(pa.nx()):
                    is_elec, pot = pa.point(x, y, z)
                    if is_elec:
                        electrode_count += 1
                        voltage_sum += pot
        
        print(f"原始电极统计:")
        print(f"  电极点数: {electrode_count}")
        print(f"  平均电压: {voltage_sum/electrode_count if electrode_count > 0 else 0:.2f} V\n")
        
        # 修改所有电极电压（例如：乘以因子）
        voltage_factor = 1.5
        print(f"将所有电极电压乘以 {voltage_factor}...")
        
        modified_count = 0
        for z in range(pa.nz()):
            for y in range(pa.ny()):
                for x in range(pa.nx()):
                    is_elec, pot = pa.point(x, y, z)
                    if is_elec:
                        new_voltage = pot * voltage_factor
                        pa.point(x, y, z, 1, new_voltage)
                        modified_count += 1
        
        print(f"已修改 {modified_count} 个电极点")
        
        # 验证修改
        new_voltage_sum = 0.0
        for z in range(pa.nz()):
            for y in range(pa.ny()):
                for x in range(pa.nx()):
                    is_elec, pot = pa.point(x, y, z)
                    if is_elec:
                        new_voltage_sum += pot
        
        print(f"\n修改后电极统计:")
        print(f"  平均电压: {new_voltage_sum/electrode_count if electrode_count > 0 else 0:.2f} V")
        
        # 保存修改后的文件
        print(f"\n保存到: {output_file}")
        pa.save(output_file)
        
        print("\n" + "="*60)
        print("修改完成!")
        print("="*60 + "\n")
        
    except Exception as e:
        print(f"\n错误: {e}\n")


def add_electrode_to_existing(input_file, output_file):
    """在现有阵列中添加新电极"""
    
    print("\n" + "="*60)
    print("添加新电极到现有阵列")
    print("="*60 + "\n")
    
    try:
        # 读取现有文件
        print(f"读取文件: {input_file}")
        pa = PA(file=input_file)
        
        print(f"阵列尺寸: {pa.nx()} x {pa.ny()} x {pa.nz()}")
        print(f"对称类型: {pa.symmetry()}\n")
        
        # 在阵列中央添加一个小电极
        center_x = pa.nx() // 2
        center_y = pa.ny() // 2
        center_z = pa.nz() // 2 if pa.nz() > 1 else 0
        
        electrode_size = 3  # 电极大小
        new_voltage = 50.0  # 新电极电压
        
        print(f"在中心添加电极:")
        print(f"  中心位置: ({center_x}, {center_y}, {center_z})")
        print(f"  电极大小: {electrode_size}x{electrode_size}")
        print(f"  电压: {new_voltage} V\n")
        
        added_count = 0
        for z in range(max(0, center_z - electrode_size//2), 
                      min(pa.nz(), center_z + electrode_size//2 + 1)):
            for y in range(max(0, center_y - electrode_size//2),
                          min(pa.ny(), center_y + electrode_size//2 + 1)):
                for x in range(max(0, center_x - electrode_size//2),
                              min(pa.nx(), center_x + electrode_size//2 + 1)):
                    pa.point(x, y, z, 1, new_voltage)
                    added_count += 1
        
        print(f"已添加 {added_count} 个电极点")
        
        # 验证添加的电极
        is_elec, pot = pa.point(center_x, center_y, center_z)
        print(f"\n验证中心点:")
        print(f"  坐标: ({center_x}, {center_y}, {center_z})")
        print(f"  电位: {pot} V")
        print(f"  是电极: {is_elec}")
        
        # 保存修改后的文件
        print(f"\n保存到: {output_file}")
        pa.save(output_file)
        
        print("\n" + "="*60)
        print("添加完成!")
        print("="*60 + "\n")
        
    except Exception as e:
        print(f"\n错误: {e}\n")


def copy_and_scale_array(input_file, output_file, scale_factor=2):
    """复制并缩放 PA 阵列"""
    
    print("\n" + "="*60)
    print(f"复制并缩放阵列 (缩放因子: {scale_factor})")
    print("="*60 + "\n")
    
    try:
        # 读取原始文件
        print(f"读取文件: {input_file}")
        pa_old = PA(file=input_file)
        
        old_nx, old_ny, old_nz = pa_old.nx(), pa_old.ny(), pa_old.nz()
        print(f"原始尺寸: {old_nx} x {old_ny} x {old_nz}")
        
        # 创建新的更大的阵列
        new_nx = old_nx * scale_factor
        new_ny = old_ny * scale_factor
        new_nz = old_nz * scale_factor if old_nz > 1 else 1
        
        print(f"新尺寸: {new_nx} x {new_ny} x {new_nz}\n")
        
        pa_new = PA(
            nx=new_nx,
            ny=new_ny,
            nz=new_nz,
            symmetry=pa_old.symmetry(),
            max_voltage=pa_old.max_voltage(),
            field_type=pa_old.field_type(),
            dx_mm=pa_old.dx_mm() / scale_factor,  # 保持物理尺寸
            dy_mm=pa_old.dy_mm() / scale_factor,
            dz_mm=pa_old.dz_mm() / scale_factor if old_nz > 1 else pa_old.dz_mm()
        )
        
        print("复制电极数据...")
        copied_count = 0
        
        # 将原始数据复制到新阵列中（放大网格）
        for z_old in range(old_nz):
            for y_old in range(old_ny):
                for x_old in range(old_nx):
                    is_elec, pot = pa_old.point(x_old, y_old, z_old)
                    
                    if is_elec:
                        # 在新阵列中填充对应的区域
                        for z_new in range(z_old * scale_factor, 
                                          (z_old + 1) * scale_factor):
                            for y_new in range(y_old * scale_factor,
                                              (y_old + 1) * scale_factor):
                                for x_new in range(x_old * scale_factor,
                                                  (x_old + 1) * scale_factor):
                                    if (x_new < new_nx and y_new < new_ny and 
                                        z_new < new_nz):
                                        pa_new.point(x_new, y_new, z_new, 1, pot)
                                        copied_count += 1
        
        print(f"已复制 {copied_count} 个电极点")
        
        # 保存新文件
        print(f"\n保存到: {output_file}")
        pa_new.save(output_file)
        
        print("\n" + "="*60)
        print("缩放完成!")
        print("="*60)
        print(f"\n物理尺寸保持不变:")
        print(f"  原始: {old_nx*pa_old.dx_mm():.2f} x {old_ny*pa_old.dy_mm():.2f} mm")
        print(f"  缩放后: {new_nx*pa_new.dx_mm():.2f} x {new_ny*pa_new.dy_mm():.2f} mm\n")
        
    except Exception as e:
        print(f"\n错误: {e}\n")


if __name__ == "__main__":
    print("\n" + "="*60)
    print("SIMION PA 文件修改示例")
    print("="*60)
    
    pa_dir = os.path.join(project_root, "SIMION", "examples", "pa")
    
    # 检查是否有可用的示例文件
    test_files = [
        "parallel_plates.pa#",
        "einzel_lens.pa#",
        "cylindrical_lens.pa#"
    ]
    
    available_files = [f for f in test_files 
                      if os.path.exists(os.path.join(pa_dir, f))]
    
    if available_files:
        print(f"\n找到可用的示例文件: {len(available_files)} 个")
        
        # 使用第一个可用文件进行演示
        input_file = os.path.join(pa_dir, available_files[0])
        
        # 示例 1: 修改电压
        output_file1 = os.path.join(pa_dir, "modified_voltage.pa#")
        modify_electrode_voltage(input_file, output_file1)
        
        # 示例 2: 添加电极
        output_file2 = os.path.join(pa_dir, "added_electrode.pa#")
        add_electrode_to_existing(input_file, output_file2)
        
        # 示例 3: 缩放阵列
        output_file3 = os.path.join(pa_dir, "scaled_array.pa#")
        copy_and_scale_array(input_file, output_file3, scale_factor=2)
        
        print("\n所有修改示例完成!")
        
    else:
        print("\n未找到可用的示例文件。")
        print("请先运行其他示例创建 PA 文件:")
        print("  python SIMION/examples/example_02_create_simple_electrode.py")
        print("\n或者提供你自己的 PA 文件路径:")
        print('  modify_electrode_voltage("你的文件.pa#", "输出文件.pa#")')

